{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "from tokenizers import BertWordPieceTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample(gpt, sentence, number=1, length=20):\n",
    "    inputs = tf.constant([tokenizer.encode(sentence).ids] * number, dtype=tf.int64)\n",
    "    length = tf.constant(length, dtype=tf.int64)\n",
    "    ret = gpt.signatures['serving_default'](inp=inputs, length=length)['output_0']\n",
    "    return [\n",
    "        tokenizer.decode(s).replace(' ', '')\n",
    "        for s in ret.numpy()\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt = hub.load('./gpt_model_tf2/')\n",
    "tokenizer = BertWordPieceTokenizer('./clue-vocab.txt',\n",
    "                                   lowercase=True,\n",
    "                                   add_special_tokens=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 36s, sys: 3.45 s, total: 1min 39s\n",
      "Wall time: 26.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ret = sample(gpt, '今天天气不错', 3, 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "今天天气不错，从机场出来已经有三个多小时了，太阳越来越高，风儿很大，想着自己要坚持续打伞避阳光，我摘下面纱帽时，有\n",
      "\n",
      "今天天气不错，格子刚刚经历了内蒙古尔班通古特大马拉松沟狂风，劈风劈雨满地飞叶子盖兔子层层翻起。牛羊一窝汗似锅头。。\n",
      "\n",
      "今天天气不错，我们出去玩了下，走了下路。看了看电影，看见了见城市。有兴趣，就带孩子们走走感觉得好的山、景、水和河，\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for s in ret:\n",
    "    print(s)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
